# $Id: trie.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

LEFT = 1
RIGHT = 2

def isLeaf(node):
   return isinstance(node, Leaf)

class Node(object):
   def __init__(self, Value = None, red = False, parent = None):
      super(Node, self).__init__()
      if Value is None:
         raise ValueError

      self._parent = parent
      self._red = red
      self._left = Leaf(self)
      self._right = Leaf(self)
      self._value = Value

   parent = property(fget = lambda self: self._parent)
   red = property(fget = lambda self: self._red)
   right = property(fget = lambda self: self._right)
   left = property(fget = lambda self: self._left)
   value = property(fget = lambda self: self._value)

   def __cmp__(self, node):
      retval = 0
      if isinstance(node, Node):
         if self.value < node.value:
            retval = -1
         elif self.value > node.value:
            retval = 1
      else:
         if node is None:
            retval = 1
         else:
            raise TypeError

      return retval

   @property
   def color(self):
      color = "black"
      if self.red:
         color = "red"

      return color

   @property
   def grandparent(self):
      retval = None
      if self.parent:
         retval = self.parent.parent

      return retval

   @property
   def sibling(self):
      retval = None
      if self.parent:
         if self is self.parent.left:
            retval = self.parent.right
         else:
            retval = self.parent.left

      return retval

   @property
   def uncle(self):
      retval = None
      if self.parent:
         retval = self.parent.sibling

      return retval

   def __repr__(self):
      return repr(str(self))

class Leaf(Node):
   def __init__(self, parent = None):
      self._parent = parent
      self._red = False
      self._left = None
      self._right = None
      self._value = None

class Tree(object):
   def __init__(self):
      super(Tree, self).__init__()
      self._root = None
      self._count = 0

   root = property(fget = lambda self: self._root)
   length = property(fget = lambda self: len(self))

   def __iter__(self):
      return self._inorder(self.root)

   def __len__(self):
      return self._count

   def _maxLeft(self, node):
      retval = None

      if isLeaf(node) or isLeaf(node.left):
         retval = node
      else:
         jm = node.left
         while not isLeaf(jm.right):
            jm = jm.right

         retval = jm

      return retval

   def _rotate(self, node, direction = 0):
      if direction == RIGHT:
         jy = node.left
         if node.parent:
            if node is node.parent.right:
               node.parent._right = jy
            else:
               node.parent._left = jy

         jy._parent = node.parent
         node._left = jy.right

         if node.left:
            node.left._parent = node

         jy._right = node
         node._parent = jy

         if not node.grandparent:
            self._root = node.parent
      elif direction == LEFT:
         jy = node.right
         if node.parent:
            if node is node.parent.right:
               node.parent._right = jy
            else:
               node.parent._left = jy

         jy._parent = node.parent
         node._right = jy.left
         if node.right:
            node.right._parent = node

         jy._left = node
         node._parent = jy

         if not node.grandparent:
            self._root = node.parent
      else:
         raise TypeError

   def _replace(self, parent, child):
      if child.parent.left is child:
         child.parent._left = Leaf(child.parent)
      else:
         child.parent._right = Leaf(child.parent)

      if not parent.parent:
         self.root = child
      elif parent.parent.left is parent:
         parent.parent._left = child
      else:
         parent.parent._right = child

      child._parent = parent.parent
      if not isLeaf(child):
         child._left = parent.left
         child.left._parent = child
         child._right = parent.right
         child.right._parent = child

   def _inorder(self, node):
      node = self.root
      while node.left:
         node = node.left

      while True:
         yield node

         if node.right:
            node = node.right
            while node.left:
               node = node.left
         else:
            while node.parent and node is node.parent.right:
               node = node.parent

            node = node.parent

         if node is None:
            break

   def insert(self, node):
      node._parent = None
      ops = 1
      if not self.root:
         self._root = node
      else:
         parent = self.root
         node._red = True
         ops = 0

         while not node.parent:
            ops = ops + 1
            if node < parent:
               if not isLeaf(parent.left):
                  parent = parent.left
                  continue

               parent._left = node
            elif node > parent:
               if not isLeaf(parent.right):
                  parent = parent.right
                  continue

               parent._right = node
            else:
               raise ValueError

            node._parent = parent

         while True:
            ip = node.parent
            if not ip:
               node._red = False
               break

            if not ip.red:
               break

            uncle = node.uncle
            if uncle.red:
               ip._red = False
               uncle._red = False
               gp = node.grandparent
               gp._red = True
               node = gp
               continue


            gp = node.grandparent
            if node is ip.right and ip is gp.left:
               self._rotate(ip, LEFT)
               node = node.left
            elif node is ip.left and ip is gp.right:
               self._rotate(ip, RIGHT)
               node = node.right

            gp = node.grandparent
            ip = node.parent
            ip._red = False
            gp._red = True

            if node is ip.left and ip is gp.left:
               self._rotate(gp, RIGHT)
            else:
               self._rotate(gp, LEFT)

            break

      self._count += 1
      return ops

   def find(self, Value):
      node = self.root
      ops = 0
      while node and node.value != Value:
         ops = ops + 1
         if Value > node.value:
            node = node.right
         else:
            node = node.left

         if isLeaf(node):
            node = None
            break

      return node

   def remove(self, node):
      if not isLeaf(node.left) and not isLeaf(node.right):
        child = self._maxLeft(node)
        node._value = child.value
        node = child

      child = node.left
      if isLeaf(child):
         child = node.right

      self._replace(node, child)
      if node.red:
         retval = 0
      elif child.red:
         child._red = False
      else:
         while True:
            parent = child.parent
            if not parent:
               break

            sibling = child.sibling
            if sibling.red:
               sibling._red = False
               parent._red = True
               if sibling is parent.left:
                  self._rotate(parent, LEFT)
               else:
                  self._rotate(parent, RIGHT)

               sibling = child.sibling

            if not parent.red and not sibling.left.red and \
                  not sibling.right.red:
               sibling._red = True
               child = parent
               continue

            if parent.red and not sibling.left.red and not sibling.right.red:
               sibling._red = True
               parent._red = False
               break

            sibling = child.sibling
            sibling._red = parent.red
            parent._red = False
            if sibling is parent.right:
               sibling.right._red = False
               self._rotate(parent, LEFT)
            else:
               sibling.left._red = False
               self._rotate(parent, RIGHT)

            break

         self._count =- 1

   def verify(self):
      result = True
      if self.root.red:
         result = False
      
      if self.root.parent:
         result = False

      bcount = None
      for node in self:
         if node.red:
            if isLeaf(node):
               result = False
               continue

            if node.left.red or node.right.red:
               result = False

         if not isLeaf(node):
            if not node.left.parent is node:
               result = False

            if not node.right.parent is node:
               result = False

         if isLeaf(node):
            tmp = node
            count = 1
            while tmp.parent:
               tmp = tmp.parent
               if not tmp.red:
                  count =+ 1

               if bcount is None:
                  bcount = count
               elif bcount != bcount:
                  result = False

      return result
